import numpy as np
import scipy
import time
import ConvertPLY
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import scipy.ndimage as ndi
import scipy.io
import sys

import numpy.linalg as la

maxiteration = 1000

Rs = np.zeros((maxiteration+1, 3, 3))
ts = np.zeros((maxiteration+1, 3))

outlier = 3
annealrate = .99

# main function for combining two point clouds
# where X is the reference cloud
# and Y is the cloud to be transformed
# R and t are initial guesses for the transformations
def probcombine(X, Y, R = np.diag([1.,1.,1.]), t = np.zeros((3,1)) ):
  global Rs, ts
  Rs[0] = R
  ts[0] = t[:,0]
  obep  = 0.
  P = np.ones((Y.shape[1], X.shape[1])) / float(Y.shape[1])
  Pi = np.ones(Y.shape[1])/Y.shape[1]
  #Praw = P/P.sum()
  sigma2 = findsigma2(X,Y,R,t)

  print 'initial %f' % objfunc(X,Y,R,t,P,sigma2,Pi)
  for i in range(maxiteration):    
    obi = objfunc(X,Y,R,t,P,sigma2,Pi)
    ######### E-step ##########
    Praw = findPraw(X,Y,R,t,sigma2,Pi)
    P =  Praw / Praw.sum(0)[None, :]
    #assert X.shape[1] + 1e3  > P.sum() > X.shape[1] - 1e3
    ####### End E-Step ########

    obe = objfunc(X,Y,R,t,P,sigma2,Pi)

    ######### M-Step ##########
    R,t = findRt(X,Y,P,sigma2)
    obm = objfunc(X,Y,R,t,P,sigma2,Pi)
    #sigma2 = findsigma2(X,Y,R,t)
    #Pi = P.sum(1)/P.sum()
    Pi = np.ones(Y.shape[1])/Y.shape[1]
    ####### End M-Step ########

    Rs[i+1,:,:] = R; ts[i+1,:] = t[:,0]
    print 'i=%d obi = %f obE = %f  obM = %f sigma2 = %f sumPi = %f'\
       % (i,obi, obe, obm, sigma2, P.sum())
    #print R
    #print t
 
    #print R.flatten()
    if obe < -600:
      break;
    obep = obe
    
    # Anneal
    sigma2 *= annealrate
    #X = rejectoutliers(X,Y,R,t)
  return R,t

# Find the variance of the mixture components
# approx: just use some sample points if set True
# use all pairs if set false
def findsigma2(X, Y, R, t, approx = False):
  assert X.shape[0] == Y.shape[0] == R.shape[0] == t.size
  Ys = np.dot(R,Y) + t
  totalerr = 0.
  totalnum = X.shape[1] * Y.shape[1]
  if approx:
    Xperm = np.random.permutation()
    raise NotImplementedError
  else:
    for i in range(Y.shape[1]):
      diff = (X - Ys[:,i:i+1])
      sqerr = (diff*diff).sum(0)
      totalerr += sqerr.sum()
    return totalerr/totalnum
  
# Find the D by R matrix P. This is the E step.
def findPraw(X,Y,R,t,sigma2,Pi):
  assert X.shape[0] == Y.shape[0] == R.shape[0] == t.size
  Ys = np.dot(R,Y) + t
  Nr = X.shape[1]; Nd = Y.shape[1]

  Praw = np.zeros((Nd, Nr))
  for r in range(Nd):
    diff = X - Ys[:,r:r+1]
    Praw[r,:] = np.exp( - (diff * diff).sum(0) / (2.*sigma2) )
  
  # return column normalized P matrix
  return Praw*Pi[:,None]
  
# Find R and t. This is the M step.
# try to find a standard optimization routine for this...
def findRt(X, Y, P, sigma2):
  return findRtdirect(X,Y,P,sigma2)

def findRtdirect(X,Y,P,sigma2):
  assert X.shape[0] == Y.shape[0]
  Nr = X.shape[1]; Nd = Y.shape[1]
  
  # Solve Nr * t = coeffX3x1 - R coeffR3x1
  # Solve R (S*Syyt - np.dot(Sy, Sy.T)) = (Sxyt * S) - np.dot(Sx, Sy.T)
  S = Nr
  tempPyt = (P.sum(1)[None,:] * Y).T
  Syyt = np.dot(Y, tempPyt)
  Sy = tempPyt.sum(0)[:,None] / S
  Sx = (P.sum(0)[None,:] * X).sum(1)[:,None] / S

  Xp = X - Sx
  Yp = Y - Sy
  
  Sxy = np.dot(np.dot(Xp, P.T), Yp.T)
  U, C, V = la.svd(Sxy)
  R = np.dot(U,V)
  t = Sx - np.dot(R, Sy)
  
  return R,t


def findRtdirect2(X,Y,P,sigma2):
  assert X.shape[0] == Y.shape[0]
  Nr = X.shape[1]; Nd = Y.shape[1]
  
  # Solve Nr * t = coeffX3x1 - R coeffR3x1
  # Solve R (S*Syyt - np.dot(Sy, Sy.T)) = (Sxyt * S) - np.dot(Sx, Sy.T)
  S = Nr
  tempPyt = (P.sum(1)[None,:] * Y).T
  Syyt = np.dot(Y, tempPyt)
  Sy = tempPyt.sum(0)[:,None]
  Sx = (P.sum(0)[None,:] * X).sum(1)[:,None]

  
  Sxyt = np.zeros((3,3))
  for r in range(Nd):
    Xscaled = P[r:r+1,:] * X
    term = (Xscaled[:,:,None] * Y[:,r][None,None,:]).sum(1)
    Sxyt += term

  A = S*Syyt - np.dot(Sy, Sy.T)
  B = Sxyt * S - np.dot(Sx, Sy.T)
  # now solve for RA = B or A'R' = B'

  soln = la.lstsq(A.T, B.T)
  R = soln[0].T; r = soln[2]
  
  t = (Sx - np.dot(R, Sy)) / S
  return R,t

def objfunc(X,Y,R,t,P,sigma2, Pi = None, bound=False):
  if bound: return upperbound(X,Y,R,t,P,sigma2)
  else: return likelihood(X,Y,R,t,P,sigma2,Pi)

def likelihood(X,Y,R,t,P,sigma2,Pi):
  llh = 0
  Ys = np.dot(R,Y) + t
  # see pg 439 in Bishop for the likelihood formula
  for c in range(X.shape[1]):
    diff = Ys - X[:,c:c+1]
    normal = np.exp(-(diff ** 2).sum(0)/(2*sigma2))/np.sqrt(sigma2)
    inside = (normal * Pi).sum()
    llh += np.log(inside)
  return llh

def upperbound(X,Y,R,t,P,sigma2):
  Nr = X.shape[1]; Nd = Y.shape[1]
  Ys = np.dot(R,Y) + t
  diffs = np.zeros((Nd,Nr))  
  for r in range(Y.shape[1]):
    diff = X - Ys[:,r:r+1]
    diffs[r,:] = (diff ** 2).sum(0)
  
  return (P * diffs).sum() / (2*sigma2)
  
def getfakedata(N=100., phi = 0, rot = .3, t = [1.,.5, .25]):
  theta = np.pi*np.arange(N)*1./N
  B = np.zeros((3,N))

  B[0,:] = 2*np.cos(theta)
  B[1,:] = np.sin(theta)
  B[2,:] = 0

  A = np.array(B)
  
  A[0,:] = 2*np.cos(theta-phi)
  A[1,:] = np.sin(theta-phi)
  A[2,:] = 0
  
  
 
  rot = np.pi/4
  rotmat = np.array([[np.cos(rot), -np.sin(rot), 0],
                        [np.sin(rot), np.cos(rot),0],
                         [0 , 0, 1.]] )
  
  A = np.dot(rotmat,A) + np.array(t)[:,None]
  return A,B

def getfakedata2(N=300., phi = 0, rot = .3, t = [-1.,.5, .25]):
  x = np.arange(N)/(N/2) - .8
  B = np.zeros((3,N))

  B[0,:] = x
  B[1,:] = x**3
  B[2,:] = 0
  
  x = np.arange(N+100)/(N/2) - 1.3
  A = np.zeros((3,N+100))

  A[0,:] = x
  A[1,:] = x**3
  A[2,:] = 0
  Aorig = np.array(A)
  A += np.random.randn(3,N+100)*0.01
  B += np.random.randn(3,N)*0.01
  
  rot = np.pi/4
  rotmat = np.array([[np.cos(rot), -np.sin(rot), 0],
                        [np.sin(rot), np.cos(rot),0],
                         [0 , 0, 1.]] )
  
  A = np.dot(rotmat,A) + np.array(t)[:,None]
  randind = (np.random.rand(N/3)*N).astype(int)
  #A = A[:,randind]
  return A,B,Aorig


def testfakedata():
  A,B,Ao = getfakedata2()
  R,t = probcombine(A,B)
  print R
  print t
  C = np.dot(R,B) + t
  import matplotlib.pyplot as plt
  plt.scatter(A[0,:],A[1,:], c='r', alpha = .5)
  plt.scatter(B[0,:],B[1,:], c='g', alpha = .5)
  plt.scatter(C[0,:],C[1,:], c='b', alpha = .5)
  plt.scatter(Ao[0,:],Ao[1,:], c='b', marker = '+',alpha = .5)
  
